use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tokio::time::sleep;

// 节点信息结构
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
struct PeerInfo {
    host: String,
    port: u16,
}

// 消息类型枚举
#[derive(Debug, Clone, Serialize, Deserialize)]
enum MessageType {
    Discovery,
    PeerList,
    Broadcast,
}

// 消息结构
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Message {
    message_type: MessageType,
    data: serde_json::Value,
}

// 节点结构
struct Node {
    host: String,
    port: u16,
    peers: Arc<Mutex<Vec<PeerInfo>>>,
    data: Arc<Mutex<HashMap<String, serde_json::Value>>>,
    shutdown_tx: Option<mpsc::Sender<()>>,
}

impl Node {
    // 创建新节点
    fn new(host: String, port: u16) -> Self {
        Node {
            host,
            port,
            peers: Arc::new(Mutex::new(Vec::new())),
            data: Arc::new(Mutex::new(HashMap::new())),
            shutdown_tx: None,
        }
    }

    // 启动节点服务器
    async fn start(&mut self) -> Result<(), Box<dyn std::error::Error>> {
        let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
        self.shutdown_tx = Some(shutdown_tx);

        let host = self.host.clone();
        let port = self.port;
        let peers = self.peers.clone();
        let data = self.data.clone();

        // 启动服务器线程
        let _server_handle = tokio::spawn(async move {
            let addr = format!("{}:{}", host, port);
            let listener = TcpListener::bind(&addr).await.expect("无法绑定到地址");
            println!("节点启动在 {}", addr);

            loop {
                tokio::select! {
                    // 监听连接
                    result = listener.accept() => {
                        match result {
                            Ok((socket, _)) => {
                                let peers_clone = peers.clone();
                                let data_clone = data.clone();
                                let host_clone = host.clone();
                                let port_clone = port;
                                
                                tokio::spawn(async move {
                                    if let Err(e) = handle_connection(
                                        socket,
                                        peers_clone,
                                        data_clone,
                                        &host_clone,
                                        port_clone,
                                    ).await {
                                        eprintln!("处理连接时出错: {:?}", e);
                                    }
                                });
                            }
                            Err(e) => {
                                eprintln!("接受连接时出错: {:?}", e);
                                break;
                            }
                        }
                    },
                    // 监听关闭信号
                    _ = shutdown_rx.recv() => {
                        println!("收到关闭信号，正在关闭服务器...");
                        break;
                    }
                }
            }
        });

        // 这里可以添加等待逻辑，或者在实际应用中保持服务器运行
        Ok(())
    }

    // 连接到已知节点
    async fn connect_to_peer(&self, peer_host: &str, peer_port: u16) -> Result<(), Box<dyn std::error::Error>> {
        let _peer_info = PeerInfo {
            host: peer_host.to_string(),
            port: peer_port,
        };

        let message = Message {
            message_type: MessageType::Discovery,
            data: serde_json::to_value(PeerInfo {
                host: self.host.clone(),
                port: self.port,
            })?,
        };

        self.send_message(peer_host, peer_port, &message).await?;
        Ok(())
    }

    // 向指定节点发送消息
    async fn send_message(&self, host: &str, port: u16, message: &Message) -> Result<(), Box<dyn std::error::Error>> {
        let mut stream = match TcpStream::connect(format!("{}:{}", host, port)).await {
            Ok(stream) => stream,
            Err(e) => {
                eprintln!("无法连接到 {}:{}: {:?}", host, port, e);
                return Err(e.into());
            }
        };

        let serialized = serde_json::to_string(message)?;
        stream.write_all(serialized.as_bytes()).await?;
        stream.write_all(b"\n").await?;
        
        Ok(())
    }

    // 向所有已知节点广播消息
    async fn broadcast(&self, data: &HashMap<String, serde_json::Value>) -> Result<(), Box<dyn std::error::Error>> {
        let message = Message {
            message_type: MessageType::Broadcast,
            data: serde_json::to_value(data.clone())?,
        };

        // 克隆节点列表以避免在迭代时锁定
        let peers = self.peers.lock().unwrap().clone();

        for peer in peers {
            if let Err(e) = self.send_message(&peer.host, peer.port, &message).await {
                eprintln!("向 {}:{} 广播消息失败: {:?}", peer.host, peer.port, e);
            }
        }

        // 同时更新自己的数据
        let mut local_data = self.data.lock().unwrap();
        for (key, value) in data {
            local_data.insert(key.clone(), value.clone());
        }

        println!("广播数据: {:?}", data);
        Ok(())
    }
}

// 处理来自其他节点的连接
async fn handle_connection(
    mut stream: TcpStream,
    peers: Arc<Mutex<Vec<PeerInfo>>>,
    data: Arc<Mutex<HashMap<String, serde_json::Value>>>,
    host: &str,
    port: u16,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut buffer = String::new();
    stream.read_to_string(&mut buffer).await?;

    // 解析消息
    let message: Message = serde_json::from_str(&buffer)?;

    match message.message_type {
        MessageType::Discovery => {
            // 处理节点发现消息
            let peer_info: PeerInfo = serde_json::from_value(message.data)?;
            
            // 检查是否需要添加新节点
            let mut should_send_peer_list = false;
            {
                let mut peers_list = peers.lock().unwrap();
                if !peers_list.contains(&peer_info) {
                    peers_list.push(peer_info.clone());
                    println!("发现新节点: {:?}", peer_info);
                    should_send_peer_list = true;
                }
                // 锁在这里自动释放
            }
            
            // 如果添加了新节点，向它发送当前节点列表
            if should_send_peer_list {
                // 创建一个新的节点实例用于发送消息
                let node = Node::new(host.to_string(), port);
                
                // 获取当前节点列表的克隆
                let current_peers = {
                    peers.lock().unwrap().clone()
                };
                
                // 构建并发送消息
                let message = Message {
                    message_type: MessageType::PeerList,
                    data: serde_json::to_value(current_peers)?,
                };
                
                if let Err(e) = node.send_message(&peer_info.host, peer_info.port, &message).await {
                    eprintln!("发送节点列表失败: {:?}", e);
                }
            }
        }
        MessageType::PeerList => {
            // 更新节点列表
            let new_peers: Vec<PeerInfo> = serde_json::from_value(message.data)?;
            let current_peer = PeerInfo { host: host.to_string(), port };
            
            let mut peers_list = peers.lock().unwrap();
            for peer in new_peers {
                if !peers_list.contains(&peer) && peer != current_peer {
                    peers_list.push(peer);
                }
            }
        }
        MessageType::Broadcast => {
            // 处理广播消息
            let broadcast_data: HashMap<String, serde_json::Value> = serde_json::from_value(message.data)?;
            println!("收到广播: {:?}", broadcast_data);
            
            // 更新本地数据
            let mut local_data = data.lock().unwrap();
            for (key, value) in broadcast_data {
                local_data.insert(key, value);
            }
        }
    }
    
    Ok(())
}

// 示例运行函数
async fn run_example() {
    // 创建并启动三个节点
    let mut node1 = Node::new("127.0.0.1".to_string(), 8080);
    let mut node2 = Node::new("127.0.0.1".to_string(), 8081);
    let mut node3 = Node::new("127.0.0.1".to_string(), 8082);
    
    // 启动节点
    node1.start().await.unwrap();
    sleep(Duration::from_millis(100)).await;
    
    node2.start().await.unwrap();
    sleep(Duration::from_millis(100)).await;
    
    node3.start().await.unwrap();
    sleep(Duration::from_millis(100)).await;
    
    // 建立节点间连接
    node2.connect_to_peer("127.0.0.1", 8080).await.unwrap();
    sleep(Duration::from_millis(100)).await;
    
    node3.connect_to_peer("127.0.0.1", 8080).await.unwrap();
    sleep(Duration::from_millis(100)).await;
    
    // 从节点1广播数据
    let mut broadcast_data = HashMap::new();
    broadcast_data.insert("key1".to_string(), serde_json::json!("value1"));
    broadcast_data.insert("key2".to_string(), serde_json::json!(42));
    
    node1.broadcast(&broadcast_data).await.unwrap();
    
    // 等待数据传播
    sleep(Duration::from_millis(100)).await;
    
    println!("演示完成");
}

#[tokio::main]
async fn main() {
    println!("分布式节点通信模拟 - Rust实现");
    println!("注意：这是一个示例程序框架，在实际运行前可能需要进一步调整");
    
    // 运行示例
    run_example().await;
}